

import abc
import base64

from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA
from Crypto.Signature import PKCS1_v1_5
from gmssl import sm2



class BaseSignatureUtil(metaclass=abc.ABCMeta):

    @abc.abstractmethod
    def signature(self, *args, **kwargs): pass

    @abc.abstractmethod
    def verify_signature(self, *args, **kwargs): pass


class SHA256WithRSAUtil(BaseSignatureUtil):

    @staticmethod
    def signature(sign_data: str, private_key: str) -> str:
        rsa_key = RSA.importKey(private_key)
        signer = PKCS1_v1_5.new(rsa_key)
        _rand_hash = SHA256.new()
        _rand_hash.update(sign_data.encode('utf-8'))
        signature = signer.sign(_rand_hash)
        return base64.b64encode(signature).decode('utf-8')

    @staticmethod
    def verify_signature(verify_data: str, signature: str, public_key: str):
        rsa_key = RSA.importKey(public_key)
        signer = PKCS1_v1_5.new(rsa_key)
        _rand_hash = SHA256.new()
        _rand_hash.update(verify_data.encode('utf-8'))
        return signer.verify(_rand_hash, base64.b64decode(signature))


class SM3WithSM2Util(BaseSignatureUtil):
    """
    SM3WithSM2签名及验签实现
    """

    @staticmethod
    def sign(sign_data: str, private_key: str, public_key: str):
    #def signature(sign_data: str, private_key: str, public_key: str):
        """
        此处签名需要公钥
        :param sign_data:
        :param private_key: 私钥
        :param public_key: 公钥
        :return: hex
        """
        #crypt = sm2.CryptSM2(private_key=private_key, public_key=public_key, mode=1, asn1=True)
        crypt = sm2.CryptSM2(private_key=private_key, public_key=public_key, mode=1, asn1=True)
        sign_str = crypt.sign_with_sm3(bytes(sign_data, encoding="utf8"))
        return sign_str
    
    @staticmethod
    def verify(verify_data: str, signature: str, public_key: str):
        """

        :param public_key:
        :param verify_data:
        :param hex signature:
        :return:
        """
        crypt = sm2.CryptSM2(private_key="", public_key=public_key, mode=1, asn1=True)
        return crypt.verify_with_sm3(signature, bytes(verify_data, encoding="utf8"))
    
    @staticmethod
    def sign_base64(sign_data: str, private_key: str, public_key: str):
    #def signature(sign_data: str, private_key: str, public_key: str):
        """
        此处签名需要公钥
        :param sign_data:
        :param private_key: 私钥
        :param public_key: 公钥
        :return: base64 str
        """
        #crypt = sm2.CryptSM2(private_key=private_key, public_key=public_key, mode=1, asn1=True)
        crypt = sm2.CryptSM2(private_key=private_key, public_key=public_key, mode=1, asn1=True)
        sign_str = crypt.sign_with_sm3(bytes(sign_data, encoding="utf8"))
        b64 = base64.b64encode(bytes.fromhex(sign_str)).decode()
        return b64

    @staticmethod
    def verify_base64(verify_data: str, signature: str, public_key: str):
        """

        :param public_key:
        :param verify_data:
        :param signature:
        :return:
        """
        crypt = sm2.CryptSM2(private_key="", public_key=public_key, mode=1, asn1=True)
        return crypt.verify_with_sm3(base64.b64decode(signature).hex(), bytes(verify_data, encoding="utf8"))
    

if __name__ == '__main__':
    rsa_pri_key = """-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCR2eBgYkRLPZgs
OFf2/81o5i69/S2rjGic/YW7Clm1mokiYZ1YLDEHH2c7tdvXJrmfLavc4ngruX4y
Uddn9A2+Ts+B5YbsbRrzOJ9ok6CA0MaVpAYejmaT/bq3e2jBY2rwI7NBzfx7c1H/
gRNY2kMTGK6l2aI4fg8jxCLt+5jOU9UbvEmwd4VNBcqfUpb9VTukdCfXt9p3jtLB
ej2bc2CQwReuv8FldhuGzDgHLvRRplXeGQ1i2OwvCguJiZH8mCq4/18336Mf0Xoc
bG7BmQYATPf+0OOkb1fJ9vDTtgLDB0T1BxiajoUNP5F4RLWer0eIDnF6iO46zwjC
czasKs3XAgMBAAECggEAEDog7ChSt9pRAW59XvZD6M1fWvT9rU3wX7MXO2Gh7IIW
itQ4eNRl2giE6FBJ9syQr2SZ7/fkbdzz7O5NHHw7QTasgKPEWK6k5nWgyrZOD7bq
uUBoC+NFB2cd7IcW9xk0DyaxrXsZAbUpZNd8CxExkEDZQYuOsGFaFt0eE8M/04i2
L9QPB3WQ3lN+sK8q9AfTBK3yv8MkGI8+00Ihp9TLZ1XqhtvlJSxKs6vLFAcH8rzP
ZKsSnqM7KGVlv5azkIeFeNLWHIzjepjA3V4ncnbu4nOBiuNGlcNh9++ELmHF+fNk
EmGO4u+GTgfuaECFtjdeq+DUH8JJ0dPXOEMhGX8p2QKBgQC5yxa3CodwW832dQ+u
pQaN3J7pZfgUrHK1idGSyGxuehEU+wyjSSwOC1eUGaP4dmj6azmvz5Ng4xPa7Ew3
TemNz7HGVnx8L+YMYa+cJHEfaWEvByin5nnPgy79Fy6aO49uK1egFGFNSPV7i4Sd
nnGt/8ye+g65uS2MwSxAN3V/WQKBgQDI9u8G0xFNI3tg7glFTr9yy/WihCZ495BZ
Bq7MAQWsaNUOR0Rh6/+gDDG2j0m3Yjnk9JFXpuctZGnVyWi1fBGjjW25e8NiyBjz
BJbpPNrkxJWpdeG51gKaOUgz2YfnooHKrt/3Cn1LO6KrNBPZMw7++NpWWh+aRE6A
cVWFOUbArwKBgHYLQdBkQS8zNQYc1CzrLuHdRZ4XKmrAMlWDTgNLkhETP10sMJhi
OjpmHGu9ar/HQ+WvUMSSRxoszIWw0q/ksbpnNpVEh+1DZY+CVVgIk1MY2iVOEBe1
SLl+qNEm9HYL15JwuUi9CiFXdJjdSRH4BFRADRsI12hK7lTauynF0sJpAoGBAIt/
fWRLVxMoshgSo7TMePCCy7tH6DzH/BcQIH3Tp9CJ5HrI1Zrzarn6PPfwdscpE+7u
JYWmgYpszVptOJXhzYyuO6ApXPNQ2qC2atr0Ny3dl9XN0iZhe+T6Jjh3o3MSxu2e
fjDJwzdZ6gzVVYHLg3lR3J1cadt31tYp1adonshZAoGAWTgi+QE1ts0HwztCw2CV
3AeKB6wZNfFNYEE4bwZhwWUQMOEjnlwh9IHwjMh5POyFD5BR/YAgihW3/u052/9O
1w7I18QvqjwSyPRVMXhItT51P5QNsugRRcsCNTHvWdozjgKuIP5pLTK+nwE7PPl3
UUNybmttjqJCyfndIOUOVlQ=
-----END PRIVATE KEY-----"""
    rsa_pub_key = """-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAkdngYGJESz2YLDhX9v/N
aOYuvf0tq4xonP2FuwpZtZqJImGdWCwxBx9nO7Xb1ya5ny2r3OJ4K7l+MlHXZ/QN
vk7PgeWG7G0a8zifaJOggNDGlaQGHo5mk/26t3towWNq8COzQc38e3NR/4ETWNpD
ExiupdmiOH4PI8Qi7fuYzlPVG7xJsHeFTQXKn1KW/VU7pHQn17fad47SwXo9m3Ng
kMEXrr/BZXYbhsw4By70UaZV3hkNYtjsLwoLiYmR/JgquP9fN9+jH9F6HGxuwZkG
AEz3/tDjpG9Xyfbw07YCwwdE9QcYmo6FDT+ReES1nq9HiA5xeojuOs8IwnM2rCrN
1wIDAQAB
-----END PUBLIC KEY-----"""
    #sm_pri_key = 'c124889fb35a23cd3092a62f92ad9f3aafd07876c5bb08bf147ac27b43e15ec7'
    #sm_pub_key = '04aaf5c364472a7b26ab254a834f5b8104ef06387ea7cc9104dd183c6ace0a647a70f19e704721919bcf955a1be8c69c9a9ead286a7d1fdfe067145244c377e1ce'
    sm_pri_key = '2AA970DCA2FE9F295E7E19794403080938C77DA6CBC2165D7F9682987D0C863'
    sm_pub_key = '6783C2C2706E0C27B82489B8E550B2D18A41DC8C7D9ED22E7B31DC8346C658B5D80111F525D23AC3A280152A40CB41214A8903832E01A6B779E0A331A4CA671F'
    
    data = "test"
    print("=========== SM3WithSM2 ===========")
    sm_sign = SM3WithSM2Util.signature(data, sm_pri_key, sm_pub_key)
    print("SM3WithSM2 计算后签名：", sm_sign)
    print("SM3WithSM2 验签结果：", SM3WithSM2Util.verify_signature(data, sm_sign, sm_pub_key))
    print("=========== SHA256WithRSA ===========")
    rsa_sign = SHA256WithRSAUtil.signature(data, rsa_pri_key)
    print("SHA256WithRSA 计算后签名：", rsa_sign)
    print("SHA256WithRSA 验签结果：", SHA256WithRSAUtil.verify_signature(data, rsa_sign, rsa_pub_key))


"""
=========== SM3WithSM2 ===========
SM3WithSM2 计算后签名： MEUCIHLCP/7f3aAdsxZHH206h65KyTJhm4ceVT1BdTjMEBZwAiEAkzRuMeD1rLQVeC0Tm7MTw/5ssEP9VWxo2dxex3WWjkQ=
SM3WithSM2 验签结果： True
=========== SHA256WithRSA ===========
SHA256WithRSA 计算后签名： FBVEU8nyALWLDKdiJFefcugGrwCcalGoOEDegbZ3N5blLHk5joIUcWHR4giFAYXm3l5wgenRNbToxItBblmGnMVkTSr6um1chGTejKtOu1KZ2koQt4hH5W8+Clc8gniuYfsLZl9W+7idl9oqczXg9vdfUqBOw/Iv1A0uEhpUdGyirf3yc2UduxKouQGGNZbVMTpOjrE4YTXs+Tbz9OIJjv19zQYcVfi1aXX4wiFJbpv/vnUGp7Bvo2FUkvLicf2r1qiPlh0ltX/EfIZmWoNvO6iD/GaQfaA9Z/7VnBhyQxhEeVMi/2KZqoVDDIOM7dIWjmCf1TGy+MDYUxzX+cGRCg==
SHA256WithRSA 验签结果： True
"""
